
import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVParser;
import org.apache.commons.csv.CSVRecord;
import org.apache.commons.math3.stat.regression.SimpleRegression;
import org.jfree.chart.ChartFactory;
import org.jfree.chart.ChartFrame;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.data.category.DefaultCategoryDataset;

import java.io.FileReader;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class 相关回归 {

    public static void main(String[] args) throws IOException {
        // 读取CSV文件
        CSVParser parser = CSVParser.parse(new FileReader("D:\\Marketing Campaign\\填充后的数据集.csv"),
                CSVFormat.DEFAULT.withFirstRecordAsHeader());

        // 定义婚姻状况映射
        Map<String, Integer> statusMap = new HashMap<>();
        statusMap.put("Single", 0);
        statusMap.put("Married", 1);
        statusMap.put("Divorced", 2);
        statusMap.put("Together", 3);
        statusMap.put("Widow", 4);
        statusMap.put("Alone", 5);
        statusMap.put("2n Cycle", 6);

        // 存储数据
        double[] incomes = new double[parser.getRecordCount()];
        int[] maritalStatuses = new int[parser.getRecordCount()];

        int index = 0;
        for (CSVRecord record : parser) {
            String maritalStatus = record.get("Marital_Status");
            if (statusMap.containsKey(maritalStatus)) {
                maritalStatuses[index] = statusMap.get(maritalStatus);
                incomes[index] = Double.parseDouble(record.get("Income"));
                index++;
            }
        }

        // 线性回归分析
        SimpleRegression regression = new SimpleRegression();
        for (int i = 0; i < index; i++) {
            regression.addData(maritalStatuses[i], incomes[i]);
        }

        // 绘制回归图
        DefaultCategoryDataset dataset = new DefaultCategoryDataset();
        for (int i = 0; i < index; i++) {
            dataset.addValue(incomes[i], "Income", String.valueOf(maritalStatuses[i]));
        }

        JFreeChart chart = ChartFactory.createLineChart(
                "婚姻状况与年收入的回归图",
                "婚姻状况",
                "年收入",
                dataset,
                PlotOrientation.VERTICAL,
                true, true, false);

        ChartFrame frame = new ChartFrame("Regression Chart", chart);
        frame.pack();
        frame.setVisible(true);
    }
}